"""
向量搜索脚本（阶段2）：
- 假设数据已通过 vector_ingest.py 入库并建索引
- 连接 Milvus，执行查询召回与相似推荐
运行： python vector_search.py
可选：设置环境变量 REQ_QUERY / REQ_CATEGORY / REQ_TOPK
"""
import os
import json
from typing import List, Dict
from pymilvus import Collection

from vector_common import (
    MILVUS_HOST, MILVUS_PORT, COLLECTION_NAME,
    connect_milvus, l2_normalize, check_tcp, OLLAMA_HOST, OLLAMA_PORT, ollama_embed
)


def search_similar(query: str, top_k: int = 5, category: str = "") -> List[Dict]:
    coll = Collection(COLLECTION_NAME)
    coll.load()
    vec = ollama_embed([query])[0]
    q = l2_normalize(vec)
    expr = f"category == '{category}'" if category else ""
    res = coll.search(
        data=[q], anns_field="embedding",
        param={"metric_type": "COSINE", "params": {"nprobe": 16}},
        limit=top_k, expr=expr,
        output_fields=["pid","name","category","brand"]
    )
    out = []
    for h in res[0]:
        out.append({
            "pid": h.entity.get("pid"),
            "name": h.entity.get("name"),
            "category": h.entity.get("category"),
            "brand": h.entity.get("brand"),
            "score": float(h.distance)
        })
    return out


def recommend_by_item(pid: int, top_k: int = 5, category: str = "") -> List[Dict]:
    coll = Collection(COLLECTION_NAME)
    coll.load()
    rows = coll.query(f"pid == {pid}", output_fields=["embedding"], limit=1)
    if not rows:
        return []
    iv = l2_normalize(rows[0]["embedding"])
    expr = f"category == '{category}'" if category else ""
    res = coll.search(
        data=[iv], anns_field="embedding",
        param={"metric_type": "COSINE", "params": {"nprobe": 16}},
        limit=top_k + 1, expr=expr,
        output_fields=["pid","name","category","brand"]
    )
    out = []
    for h in res[0]:
        pid2 = h.entity.get("pid")
        if pid2 == pid:
            continue
        out.append({
            "pid": pid2,
            "name": h.entity.get("name"),
            "category": h.entity.get("category"),
            "brand": h.entity.get("brand"),
            "score": float(h.distance)
        })
        if len(out) >= top_k:
            break
    return out


def main():
    assert check_tcp(OLLAMA_HOST, OLLAMA_PORT), "Ollama not reachable"
    assert check_tcp(MILVUS_HOST, int(MILVUS_PORT)), "Milvus not reachable"
    connect_milvus()

    query = os.getenv("REQ_QUERY", "所有美妆商品")
    category = os.getenv("REQ_CATEGORY", "")
    top_k = int(os.getenv("REQ_TOPK", "5"))

    sims = search_similar(query, top_k=top_k, category=category)
    print("\n搜索结果:")
    print(json.dumps(sims, ensure_ascii=False, indent=2))

    recs = recommend_by_item(1, top_k=top_k, category=category)
    print(f"\n基于商品 1 的相似推荐:")
    print(json.dumps(recs, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
